# modified from https://github.com/desimone/vision/blob/fb74c76d09bcc2594159613d5bdadd7d4697bb11/torchvision/datasets/folder.py

import os
import os.path

import torch
from torchvision import transforms
import torch.utils.data as data
from PIL import Image

def load_list_from_file(file_name):
    """
    从文本文件中加载内容到列表中。

    参数:
        file_name (str): 文本文件的路径。

    返回:
        list: 文件中的内容，每行作为一个元素。
    """
    try:
        with open(file_name, "r", encoding="utf-8") as file:
            # 逐行读取文件内容，并去除每行的首尾空白字符
            loaded_list = [line.strip() for line in file.readlines()]
        return loaded_list
    except FileNotFoundError:
        print(f"文件 {file_name} 未找到，请检查路径是否正确！")
        return []
    except Exception as e:
        print(f"读取文件时发生错误：{e}")
        return []


def default_loader(path):
    return Image.open(path).convert('RGB')

class ImageFolder(data.Dataset):
    """ ImageFolder can be used to load images where there are no labels."""

    def __init__(self, set_path, transform=None, loader=default_loader):
        self.imgs = load_list_from_file(set_path)
        self.transform = transform
        self.loader = loader

    def __getitem__(self, index):
        filepath = self.imgs[index]
        try:
            img = self.loader(filepath)
        except:
            return torch.zeros((3, 32, 32))

        if self.transform is not None:
            try:
                img = self.transform(img)
            except:
                return torch.zeros((3, 32, 32))
        return img

    def __len__(self):
        return len(self.imgs)
    